#!/usr/bin/python3
# -*- coding: utf-8 -*-

import random, os, sys, time, asyncio, threading, logging, re
from io import BytesIO
from urllib import request
from urllib.error import URLError, HTTPError
from contextlib import closing

import aiohttp, requests
from Crypto.Cipher import AES

from logger import set_logger

success = {}

def random_ip():
    a=random.randint(1,255)
    b=random.randint(1,255)
    c=random.randint(1,255)
    d=random.randint(1,255)
    return '%s.%s.%s.%s' % (str(a), str(b), str(c), str(d))

def random_Agent():
    agents = [
    "Mozilla/5.0 (Windows NT 6.3; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/35.0.1916.153 Safari/537.36",
    "Mozilla/5.0 (Windows NT 6.1; WOW64; rv:30.0) Gecko/20100101 Firefox/30.0",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_2) AppleWebKit/537.75.14 (KHTML, like Gecko) Version/7.0.3 Safari/537.75.14",
    "Mozilla/5.0 (compatible; MSIE 10.0; Windows NT 6.2; Win64; x64; Trident/6.0)",
    'Mozilla/5.0 (Windows; U; Windows NT 5.1; it; rv:1.8.1.11) Gecko/20071127 Firefox/2.0.0.11',
    'Opera/9.25 (Windows NT 5.1; U; en)',
    'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1; SV1; .NET CLR 1.1.4322; .NET CLR 2.0.50727)',
    'Mozilla/5.0 (compatible; Konqueror/3.5; Linux) KHTML/3.5.5 (like Gecko) (Kubuntu)',
    'Mozilla/5.0 (X11; U; Linux i686; en-US; rv:1.8.0.12) Gecko/20070731 Ubuntu/dapper-security Firefox/1.5.0.12',
    'Lynx/2.8.5rel.1 libwww-FM/2.14 SSL-MM/1.4.1 GNUTLS/1.2.9',
    "Mozilla/5.0 (X11; Linux i686) AppleWebKit/535.7 (KHTML, like Gecko) Ubuntu/11.04 Chromium/16.0.912.77 Chrome/16.0.912.77 Safari/535.7",
    "Mozilla/5.0 (X11; Ubuntu; Linux i686; rv:10.0) Gecko/20100101 Firefox/10.0 ",
    'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:73.0) Gecko/20100101 Firefox/73.0',
    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/70.0.3538.110 Safari/537.36'
    ]
    return random.choice(agents)

def add_header(text=True):
    base = {
        "Accept-Language": "zh-CN,zh;q=0.9,en-US;q=0.8,en;q=0.7", 
        "Proxy-Connection": "keep-alive", 
        'User-Agent': random_Agent(), 
        'X-Forwarded-For':random_ip(),
    }
    if not text:
        base['Accept-Encoding'] = 'gzip, deflate, br'
    return base

def download(url, filename):
    result = None
    error_msg = {}
    try:
        logging.info('[-]access: %s' % url)
        req = request.Request(url)
        for key, value in add_header().items():
            req.add_header(key, value)
        with request.urlopen(req) as resp:
            if resp.status == 200:
                with open(filename, 'ab') as dst:
                    dst.write(resp.read())
            else:
                logging.warning('[-]failed to access %s: %s, %s' % (url, resp.status, resp.reason))
        if os.path.exists(filename):
            logging.info('[-]success to download [%s]' % filename)
        else:
            logging.warning('[-]failed to download [%s]' % filename)
    except HTTPError as e:
        error_msg = ({url: 'Error: %s' % e})
    except URLError as e:
        error_msg = ({url: 'Error: %s' % e})
    except BaseException as e:
        error_msg= ({url: 'Error: %s' % e})
    if len(error_msg) > 0:
        logging.error('[-]%s: %s' % (threading.current_thread().getName(), error_msg))

async def downloadfile(url, filename, params, key=None, chunk_size = 1024, bad_req=set(),run_key={}, text=False):
    headers = add_header(text)
    result = None
    error_msg = {}
    try:
        with closing(requests.get(url, allow_redirects=True, timeout=(30,30), headers=add_header(text), params=params, stream=True)) as response:
            if not run_key.get(0, True): raise BaseException('run_key exit: %s' % run_key)
            if response.status_code == 200:
                buf = BytesIO()
                for data in response.iter_content(chunk_size=chunk_size):
                    if not run_key.get(0, True): raise BaseException('run_key exit: %s' % run_key)
                    buf.write(data)
                buf.seek(0)
                result = buf.read()
                buf.seek(0)
                result = buf.read()
            else:
                bad_req.add(url)
                error_msg = {url:  'params: %s, status: %s, reason: %s]\n%s' % (params, response.status_code, response.reason, response.content)}
            if result:
                with open(filename, 'ab') as file:
                    # if key and len(key) > 0:
                    #     cryptor = AES.new(key, AES.MODE_CBC, key)
                    #     result = cryptor.decrypt(result)
                    file.write(result)
            else:
                bad_req.add(url)
                if not error_msg:
                    error_msg = ({url: 'result is %s' % result})
    except BaseException as e:
        error_msg = ({url: 'Error: %s' % e})
    if len(error_msg) > 0:
        logging.info('[-]%s: %s' % (threading.current_thread().getName(), error_msg))

async def downloadfile_(url, filename, params, key=None, chunk_size = 1024, bad_req=set(),run_key={}, text=False):
    headers = add_header(text)
    result = None
    error_msg = {}
    try:
        async with aiohttp.ClientSession(headers=headers) as session:
            if not run_key.get(0, True): raise BaseException('run_key exit: %s' % run_key)
            async with session.get(url, params=params) as response:
                if response.status == 200:
                    buf = BytesIO()
                    while True:
                        if not run_key.get(0, True): raise BaseException('run_key exit: %s' % run_key)
                        chunk = await response.content.read(chunk_size)
                        if not chunk:
                            break
                        buf.write(chunk)
                    buf.seek(0)
                    result = buf.read()
                else:
                    bad_req.add(url)
                    error_msg = {url:  'params: %s, status: %s, reason: %s]\n%s' % (params, response.status, response.reason, response.content)}
            if result:
                with open(filename, 'ab') as file:
                    # if key and len(key) > 0:
                    #     cryptor = AES.new(key, AES.MODE_CBC, key)
                    #     result = cryptor.decrypt(result)
                    file.write(result)
            else:
                bad_req.add(url)
                if not error_msg:
                    error_msg = ({url: 'result is %s' % result})
    except BaseException as e:
        error_msg = ({url: 'Error: %s' % e})
    if len(error_msg) > 0:
        logging.info('[-]%s: %s' % (threading.current_thread().getName(), error_msg))

def resource_name(url):
    name = url.rsplit('/', 1)[1]
    return name if '?' not in name else name.split('?', 1)[0]

def downloadtask(urls, savepath, retry=5, key=None, bad_req = set(), run_key={}, text=False):
    tasks = []
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        os.makedirs(savepath, exist_ok = True)
        urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url)))]
        count = 0
        while(len(urls) > 0):
            logging.info('[-]task of %s: %s' % (threading.current_thread().getName(), len(urls)))
            if not run_key.get(0, True): break
            tasks = [downloadfile(url, '%s/%s' % (savepath, resource_name(url)), {}, key=key, bad_req=bad_req, run_key=run_key, text=text) for url in urls]
            try:
                loop.run_until_complete(asyncio.wait(tasks))
            except KeyboardInterrupt as e:
                logging.info(asyncio.Task.all_tasks())
                for task in asyncio.Task.all_tasks():
                    logging.info(task.cancel())
                loop.stop()
                # loop.run_forever()
                logging.error('[-]Error: %s' % e)
                break
            except BaseException as e:
                logging.error('[-]Error: %s' % e)
                break
            urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url))) and url not in bad_req]
            count += 1
            if count >= retry: break
    except BaseException as e:
        logging.error('Error: %s' % e)
    finally:
        loop.close()

def checktask_status(task_tss, savepath, bad_req = set(), run_key={}):
    tmp_file = '%s/status.tmp' % (savepath)
    with open(tmp_file, 'w') as status_file:
        status_file.write('[-]task start...\n')
    while True:
        try:
            flag = True
            status = ''
            sum = 0
            for index, urls in enumerate(task_tss):
                urls = [url for url in urls if not os.path.exists('%s/%s' % (savepath, resource_name(url))) and url not in bad_req]
                status += '%6d\t' % len(urls)
                sum += len(urls)
                if len(urls) > 0: flag = False
            if not os.path.exists(tmp_file): 
                success[threading.current_thread().name] = False
                run_key[0] = False
                break
            with open(tmp_file, 'r+') as status_file:
                status_file.seek(0, 0)
                status_file.write('\n[-]task remain: %s\t | sum(%s)' % (status, sum))
            if flag: break
        except KeyboardInterrupt as e:
            break;
    logging.info('download task complete')

def getm3u8key(key_url, savepath):
    if not key_url or len(key_url) <= 0: return
    with request.urlopen(key_url) as f:
        if f.status == 200:
            val = f.read()
            filepath = savepath
            os.makedirs(filepath.rsplit('/', 1)[0], exist_ok = True)
            with open(filepath, 'wb') as dst:
                dst.write(val)
            return val
        else:
            logging.info('[-]Status:', f.status, f.reason)

def getm3u8key_(baseurl, keys, savepath):
    if len(keys) <= 0: return
    beg_m = keys.index('METHOD=') + 7
    end_m = keys.index(',')
    beg_u = keys.index('URI=') + 4
    method = keys[beg_m:end_m]
    if keys.index('"') > 0 or keys.index("'") > 0:
        key = keys[beg_u+1:-1]
    else:
        key = keys[beg_u:]
    # url_root = '%s%s' % (baseurl[:baseurl.index('://') + 3], baseurl[baseurl.index('://') + 3:].split('/', 1)[0])
    url = '%s/%s' % (baseurl, key)
    with request.urlopen(url) as f:
        if f.status == 200:
            val = f.read()
            filepath = '%s/%s' % (savepath, key)
            os.makedirs(filepath.rsplit('/', 1)[0], exist_ok = True)
            with open(filepath, 'wb') as dst:
                dst.write(val)
            return val
        else:
            logging.info('[-]Status:', f.status, f.reason)

def downloadm3u8(url, savepath, result = {}, filename=None):
    if ':' in savepath: raise BaseException('":" in savepath')
    os.makedirs(savepath, exist_ok = True)
    all_content = ''
    cur_file_name = resource_name(url)
    filepath = '%s/%s' % (savepath, cur_file_name)
    if os.path.exists(filepath):
        logging.warning('remove m3u8 file: %s' % filepath)
        os.remove(filepath)
    logging.info('[-]download m3u8 file: %s' % url)
    downloadtask([url], savepath)
    if not os.path.exists(filepath):
        raise BaseException("m3u8 file not exists: %s" % filepath)
    with open(filepath, 'r') as m3u8:
        all_content = m3u8.read()
    if len(all_content) <= 0:
        raise BaseException("Failed to download m3u8 file: \n%s" % all_content)
    if '\r\n' in all_content:
        file_line = all_content.split("\r\n")
    else:
        file_line = all_content.split("\n")
    if file_line[0] != "#EXTM3U":
        raise BaseException("it is not m3u8 file: \n%s" % all_content)
    tss = []
    keys = ''
    url , name = url.rsplit('/', 1)
    url_root = '%s%s' % (url[:url.index('://') + 3], url[url.index('://') + 3:].split('/', 1)[0])
    EXTURL_flag = False
    ROOT_flag = False
    duration = 0.0
    key_url = None
    key_path = None
    with open('%s/%s' % (savepath, cur_file_name), 'w') as file:
        for line in file_line:
            if line.startswith('/'):
                line_ = line[1:]
                ROOT_flag = True
            else:
                line_ = line
            # line_ = line[1:] if line.startswith('/') else line
            if '#EXT-X-KEY' in line:
                keys = line
                key_uri = keys[keys.index('URI=') + 4:]
                key_uri = key_uri[1:-1] if '"' in key_uri or "'" in key_uri else key_uri
                if key_uri.startswith('/'):
                    key_uri = key_uri[1:]
                    key_url = '%s/%s' % (url_root, key_uri)
                else:
                    key_uri = key_uri
                    key_url = '%s/%s' % (url, key_uri)
                key_path = '%s/%s' % (savepath, key_uri)
                keys = '%sURI="%s"' % (keys[:keys.index('URI=')], key_uri)
                line_ = keys
            if '#EXTURL:' in line:
                url = line[8:]
                EXTURL_flag = True
            if '#EXTINF:' in line:
                exp = re.match('#EXTINF:(.\d*\.?\d*),', line)
                duration = duration + float(exp.group(1))
            file.write('%s\n' % line_)
            if "#EXT" in line or len(line.strip()) == 0: continue
            tss.append(line_)
        if not EXTURL_flag:
            if ROOT_flag:
                EXTURL = url_root
            else:
                EXTURL = url
            file.write('\n#EXTURL:%s' % EXTURL)
    tss_out = []
    for ts in tss:
        # print(ts)
        cur_url = ''
        if ROOT_flag:
            cur_url = '%s/%s' % (url_root, ts)
        else:
            cur_url = '%s/%s' % (url, ts)
        ts_path = ts.rsplit('/', 1)[0] if '/' in ts else None
        ts_name = ts.replace('/', '_') if '/' in ts else ts
        if ts_path:
            savepath_ = '%s/%s' % (savepath, ts_path)
        else:
            savepath_ = savepath
        if '.m3u8' in ts:
            downloadm3u8(cur_url, savepath_, result, ts_name)
        else:
            tss_out.append(cur_url)
    if len(tss_out) > 0:
        key = getm3u8key(key_url, key_path)
        filename = (filename if '.' not in filename else filename.rsplit('.', 1)[0]) if filename else name
        result[filename] = {'tss': tss_out, 'key': key, 'url': url, 'path': savepath_, 'duration': duration}

def loadm3u8file(savepath, filename, result):
    all_content = ''
    filepath = '%s/%s' % (savepath, filename)
    logging.info('[-]loading m3u8 file: %s' % (filepath))
    if not os.path.exists(filepath): raise BaseException('file not exists: %s' % filepath)
    with open(filepath, 'r') as file:
        all_content = file.read()
    if '\r\n' in all_content:
        file_line = all_content.split("\r\n")
    else:
        file_line = all_content.split("\n")
    if file_line[0] != "#EXTM3U":
        raise BaseException("it is not m3u8 file: \n%s" % all_content)
    tss = []
    keys = ''
    duration = 0.0
    url = None
    for line in file_line:
        if '#EXT-X-KEY' in line:
            keys = line
        if '#EXTURL:' in line:
            url = line[8:]
        if '#EXTINF:' in line:
            exp = re.match('#EXTINF:(.\d*\.?\d*),', line)
            duration = duration + float(exp.group(1))
        if "#EXT" in line or len(line.strip()) == 0: continue
        tss.append(line)
    if not url: raise BaseException('url is None: %s' % url)
    tss_out = []
    for ts in tss:
        # print(ts)
        cur_url = '%s/%s' % (url, ts)
        if '.m3u8' in ts:
            ts_path = ts.rsplit('/', 1)[0] if '/' in ts else None
            ts_name = ts.replace('/', '_') if '/' in ts else ts
            if ts_path:
                savepath_ = '%s/%s' % (savepath, ts_path)
            else:
                savepath_ = savepath
            downloadm3u8('%s/%s' % (url, ts), savepath_, result, ts_name)
        else:
            tss_out.append(cur_url)
    if len(tss_out) > 0:
        key = getm3u8key_(url, keys, savepath)
        filename = filename if '.' not in filename else filename.rsplit('.', 1)[0]
        result[filename] = {'tss': tss_out, 'key': key, 'url': url, 'path': savepath, 'duration': duration}
    
def download_batch_ts(tss, savepath, key=None, run_key = {}):
    size = len(tss)
    pagesize = size // 10 + 1
    bad_req = set()
    run_key[0] = True
    for i in range(0, size, pagesize):
        t = threading.Thread(target=downloadtask, name='%s-%s' % (downloadtask.__name__, i // pagesize), args=(tss[i:i+pagesize], savepath), kwargs={'key': key, 'bad_req': bad_req, 'run_key': run_key})
        t.start()
    t = threading.Thread(target=checktask_status, name='%s-%s' % (checktask_status.__name__, 1), args=([tss[i:i+pagesize] for i in range(0, size, pagesize)], savepath), kwargs={'bad_req': bad_req, 'run_key': run_key})
    t.start()
    t.join()

def check_files(tss, path):
    ts_files = ['%s/%s/%s' % (os.getcwd(), path, resource_name(ts)) for ts in tss]
    file_not_exists = []
    exists = 0
    not_exists = 0
    for index, file_ in enumerate(ts_files):
        if not os.path.exists(file_):
            file_not_exists.append(tss[index])
            not_exists += 1
        else:
            exists += 1
    logging.info('[-]check_files in %s: [total: %-8d| exists: %-8d| not_exists: %-8d]' % (path, len(ts_files), exists, not_exists))
    return file_not_exists

def mergefiles(tss, path, filename, key, rm_src = False, cover = True, ignore=True, run_key = {}):
    ts_files = ['%s/%s/%s' % (os.getcwd(), path, resource_name(ts)) for ts in tss]
    if len(ts_files) == 0: raise BaseException('num of ts_files is 0: %s' % ts_files)
    if len(filename) == 0: raise BaseException('filename is null: %s' % filename)
    file_not_exists = []
    file_exists = []
    for file_ in ts_files:
        if not os.path.exists(file_):
            file_not_exists.append(file_)
        else:
            file_exists.append(file_)
    if len(file_not_exists) > 0 and not ignore:
        raise BaseException('these file is not exists: %s' % file_not_exists)
    filename = '%s/%s' % (path, filename)
    if os.path.exists(filename):
        if cover:
            os.remove(filename)
            logging.info('[-]cover dst file %s' % filename)
        else:
            filename = '%s_%s' % (filename, time.time())
    logging.info('[-]dst file save as: %s' % filename)
    with open(filename, 'ab') as dst:
        filesize = 0
        for index, file_ in enumerate(file_exists):
            with open(file_, 'rb') as src:
                result = src.read()
                if key and len(key) > 0:
                    cryptor = AES.new(key, AES.MODE_CBC, key)
                    result = cryptor.decrypt(result)
                filesize = filesize + len(result)
                dst.write(result)
        run_key['size'] = filesize
        #     print('[-]%s: %s in %s | %3.2f%%' % ('mergefiles', index + 1, len(file_exists), ((index + 1) / len(file_exists) * 100.0)), end = '\r')
        # print('')
        logging.info('[-]mergefiles %s files to %s' % (len(file_exists), filename))
    if rm_src:
        logging.info('[-]remove src files')
        for file_ in file_exists:
            os.remove(file_)

def main_file(path, m3u8, info=False):
    # 文件夹生成同名文件
    filename = path
    result = {}
    loadm3u8file(path, m3u8, result)
    if info:
        logging.info(result)
        return
    download_ts(result)

def main_url(url, filename, info=False):
    # 生成同名文件夹
    path = filename
    result = {}
    downloadm3u8(url, path, result)
    if info:
        logging.info(result)
        return
    download_ts(result)

def download_ts(result, run_key = {}):
    if result is None or len(result) <= 0: return
    for k, v in result.items():
        filename = k
        tss = v['tss']
        key = v['key']
        url = v['url']
        path = v['path']
        tss_not_exists = check_files([ts for ts in tss], path)
        download_batch_ts([ts for ts in tss_not_exists], path, key=key, run_key = run_key)
        mergefiles([ts for ts in tss], path, '%s.mp4' % filename, key, run_key = run_key)

if __name__ == '__main__':
    help_msg = '''\
        Usage:python3 %s -u url -p path -f filename
        -u: download url
        -p: dst save as <path>.mp4
        -f: load m3u8 file <path/filename>
        require pycryptodome
        pip3 install pycryptodome -i https://pypi.tuna.tsinghua.edu.cn/simple
        '''    % sys.argv[0]
    # print(sys.argv)
    if len(sys.argv) < 2:
        print(help_msg)
    else:
        start_t = time.time()
        set_logger()
        # main_file
        if '-f' in sys.argv:
            filename = sys.argv[sys.argv.index('-f') + 1]
            if '-p' in sys.argv:
                path = sys.argv[sys.argv.index('-p') + 1]
            else:
                path, filename = filename.rsplit('/', 1)
            info = False
            if '--info' in sys.argv:
                info = True
            m3u8 = '%s/%s' % (path, filename)
            if not os.path.exists(m3u8):
                raise BaseException('file not exists: %s' % m3u8)
            main_file(path, filename, info)
            spend = int(time.time() - start_t)
            logging.info('[-]spend: %ds' % (spend))
            exit(0)
        # main_url
        if '-u' in sys.argv:
            url = sys.argv[sys.argv.index('-u') + 1]
            if '-p' not in sys.argv:
                path = resource_name(url)
                if '.' in path:
                    path = path.rsplit('.', 1)[0]
            else:
                path = sys.argv[sys.argv.index('-p') + 1]
            info = False
            if '--info' in sys.argv:
                info = True
            logging.info('[DEBUG] path: %s' % path)
            main_url(url, path, info)
            spend = int(time.time() - start_t)
            logging.info('[-]spend: %ds' % (spend))
            exit(0)
        if '--test' in sys.argv:
            url = sys.argv[sys.argv.index('--test') + 1]
            name = resource_name(url)
            logging.info('[-]%s' % name)
            exit(0)
        if '--trans' in sys.argv:
            exit(0)
        print(help_msg)
